""""""
import os
import torch
import numpy as np
import onnxruntime
from tqdm import tqdm
from PIL import Image
from .util import Util
from insightface.app import FaceAnalysis


class LoadImagesFromFolder:
    """
    从文件夹中批量加载图像
    """
    
    CATEGORY = "destiny/图像"

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "dirname": ("STRING", {
                    "default": "input",
                    "placeholder": "请输入加载路径"
                }),
                "include_png": ("BOOLEAN", {"default": True}),
                "include_jpg_jpeg": ("BOOLEAN", {"default": True}),
                "mode": (["剪裁", "缩放", "填充"], {"default": "缩放"}),
                "width": ("INT", {"default": 512, "min": 64, "max": 4096}),
                "height": ("INT", {"default": 512, "min": 64, "max": 4096}),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING", "STRING")
    RETURN_NAMES = ("images", "file_names", "dirname")

    FUNCTION = "execute"

    def execute(self, dirname, include_png=False, include_jpg_jpeg=False, mode="缩放", width=512, height=512):
        dirname = os.path.abspath(dirname)
        assert os.path.exists(dirname), f"[destiny] 文件夹不存在: {dirname}"
        assert os.path.isdir(dirname), f"[destiny] 非有效文件夹: {dirname}"
        include_exts = []
        if include_png:
            include_exts.append('.png')
        if include_jpg_jpeg:
            include_exts.extend(['.jpg', '.jpeg'])
        tensor_images = []
        out_file_names = []
        t = tqdm(os.listdir(dirname))
        for ix, file_name in enumerate(t):
            _, f_ext = os.path.splitext(file_name)
            if str(f_ext).lower() in include_exts:
                t.set_description(f"[Deal {ix+1}/{len(t)}]")
                t.set_postfix(file_name=file_name)
                out_file_names.append(file_name)
                filename = os.path.join(dirname, file_name)
                image = Image.open(filename).convert('RGB')
                if mode == '剪裁':
                    i_w, i_h = image.size
                    if i_w > width: 
                        i_w, half_w = int(i_w/2), int(width/2)
                        o_x1, o_x2 = i_w - half_w, i_w + half_w
                        # print(f"[destiny] crop:", (o_x1, 0, o_x2, i_h))
                        image = image.crop((o_x1, 0, o_x2, i_h))
                        i_w, i_h = image.size
                    if i_h > height: 
                        i_h, half_h = int(i_h/2), int(height/2)
                        o_y1, o_y2 = i_h - half_h, i_h + half_h
                        # print(f"[destiny] crop:", (0, o_y1, i_w, o_y2))
                        image = image.crop((0, o_y1, i_w, o_y2))
                elif mode == '填充':
                    np_img = np.array(image)
                    np_img_height, np_img_width, _c = np_img.shape
                    pad_total = abs(np_img_height-np_img_width)
                    pad_1 = int(pad_total/2)
                    pad_2 = pad_total - pad_1
                    if np_img_height > np_img_width:
                        np_img = np.pad(np_img, ((0,0), (pad_1, pad_2), (0,0)), mode='constant', constant_values=[255])
                    elif np_img_width > np_img_height:
                        np_img = np.pad(np_img, ((pad_1, pad_2), (0, 0), (0,0)), mode='constant', constant_values=[255])
                    image = Image.fromarray(np_img)
                image = image.resize((width, height))
                tensor_images.append(Util.pil2tensor(image))
        assert len(tensor_images) > 0, f"[destiny] 图片数量不能为空！"
        tensor_images = torch.stack(tensor_images, dim=0)
        return (tensor_images, out_file_names, dirname)
    

class SaveImages:
    """
    保存图像
    """
    
    CATEGORY = "destiny/图像"

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE", {}),
                "file_names": ("STRING", {"forceInput": True}),
            },
            "optional": {
                "exists_ok": ("BOOLEAN", {"default": False}),
                "save_path": ("STRING", {"default": "output"}),
                "prefix": ("STRING", {"default": ""}),
            }
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True

    FUNCTION = "execute"

    def execute(self, images, file_names: str, exists_ok=False, save_path="output", prefix=""):
        save_path = os.path.abspath(save_path)
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        for image, file_name in zip(images, file_names):
            img = Util.tensor2pil(image)
            target_filename = os.path.join(save_path, f"{prefix}{file_name}")
            if not os.path.exists(target_filename) or exists_ok:
                img.save(target_filename)
            # print(save_path, file_name)
        return ()


class ImageInfo:
    """
    获取图像信息
    """
    
    CATEGORY = "destiny/图像"

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE", {})
            },
        }

    RETURN_TYPES = ("INT", "INT", "INT")
    RETURN_NAMES = ("h", "w", "c")
    OUTPUT_NODE = True

    FUNCTION = "execute"

    def execute(self, image):
        print('image', type(image))
        return (200, 300, 3)


class DetectFace:
    """
    面部检测
    """
    
    CATEGORY = "destiny/图像"

    def __init__(self):
        self.model_path = os.path.join(
            os.getcwd(), 'models', 'insightface'
        )

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE", {}),
                "det_thresh": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1, "step": 0.01}),
                "draw_kps": ("BOOLEAN", {"default": True}),
                "scale_width": ("FLOAT", {"default": 1, "min": 0.1, "step": 0.01}),
                "scale_height": ("FLOAT", {"default": 1, "min": 0.1, "step": 0.01}),
                "width":("INT", {"default": 512}),
                "height":("INT", {"default": 512}),
                "mode": (["填充", "拉伸"], {"default": "填充"})
            },
        }

    RETURN_TYPES = ("IMAGE", "IMAGE")
    RETURN_NAMES = ("label_image", "face_image")
    OUTPUT_NODE = True

    FUNCTION = "execute"

    def get_model(self, det_thresh=0.5):
        providers = onnxruntime.get_available_providers()
        model = FaceAnalysis(
            name='buffalo_l',
            root=self.model_path,
            allowed_modules=('detection', ),
            providers=providers,
        )
        model.prepare(ctx_id=0, det_thresh=det_thresh)
        return model

    def execute(self, image, det_thresh=0.5, scale_width=1, scale_height=1, draw_kps=False, width=512, height=512, mode='填充'):
        outputs = []
        face_imgs = []
        scores = []
        # bboxes = []
        model = self.get_model(det_thresh=det_thresh)
        for batch_ix, img in enumerate(image):
            img = Util.tensor2pil(img)
            img_width, img_height = img.size
            results = []
            faces = model.get(np.array(img))
            for face in faces:
                bbox = np.array(face.bbox).astype(np.int32).tolist()
                bbox = Util.scale_bbox(bbox=bbox, scale_width=scale_width, scale_height=scale_height)
                if bbox[-2] > img_width:
                    bbox[-2] = img_width
                if bbox[-1] > img_height:
                    bbox[-1] = img_height
                result = {
                    'bbox': bbox,
                    'kps': np.array(face.kps).astype(np.int32).tolist(),
                    'det_score': face.det_score.item(),
                }
                results.append(result)
                face_img = img.crop(tuple(bbox))
                if width == height and mode == '填充':
                    np_img = np.array(face_img)
                    np_img_height, np_img_width, _c = np_img.shape
                    pad_total = abs(np_img_height-np_img_width)
                    pad_1 = int(pad_total/2)
                    pad_2 = pad_total - pad_1
                    if np_img_height > np_img_width:
                        np_img = np.pad(np_img, ((0,0), (pad_1, pad_2), (0,0)), mode='constant', constant_values=[0])
                    elif np_img_width > np_img_height:
                        np_img = np.pad(np_img, ((pad_1, pad_2), (0, 0), (0,0)), mode='constant', constant_values=[0])
                    face_img = Image.fromarray(np_img)
                face_img = face_img.resize((width, height))
                face_imgs.append(Util.pil2tensor(face_img))
            if len(results) > 0:
                v_img = Util.draw(img, results, draw_kps=draw_kps)
                outputs.append(Util.pil2tensor(v_img))
        assert len(outputs) > 0 and len(face_imgs) > 0, f"[destiny] 未检测到人脸"
        return (torch.stack(outputs, dim=0), torch.stack(face_imgs, dim=0))
        